#用法
#python dcgan_fashion_mnist.py-输出输出

#导入必要的软件包
from pyimagesearch.dcgan import DCGAN
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.datasets import fashion_mnist
from sklearn.utils import shuffle
from imutils import build_montages
import numpy as np
import argparse
import cv2
import os

# 构造参数解析并解析参数
ap = argparse.ArgumentParser()
ap.add_argument("-o", "--output", required=True,
	help="path to output directory")
ap.add_argument("-e", "--epochs", type=int, default=50,
	help="# epochs to train for")
ap.add_argument("-b", "--batch-size", type=int, default=128,
	help="batch size for training")
args = vars(ap.parse_args())

# 将时期和批量大小存储在便利变量中，然后
# 初始化我们的学习率
NUM_EPOCHS = args["epochs"]
BATCH_SIZE = args["batch_size"]
INIT_LR = 2e-4

# 加载Fashion MNIST数据集并堆叠训练和测试
# 数据点，所以我们有其他训练数据
print("[INFO] loading MNIST dataset...")
((trainX, _), (testX, _)) = fashion_mnist.load_data()
trainImages = np.concatenate([trainX, testX])

# 为通道添加额外的尺寸并缩放图像
# 进入[-1，1]范围（正切范围
# 功能）
trainImages = np.expand_dims(trainImages, axis=-1)
trainImages = (trainImages.astype("float") - 127.5) / 127.5

# 建立生成器
print("[INFO] building generator...")
gen = DCGAN.build_generator(7, 64, channels=1)

# 建立鉴别器
print("[INFO] building discriminator...")
disc = DCGAN.build_discriminator(28, 28, 1)
discOpt = Adam(lr=INIT_LR, beta_1=0.5, decay=INIT_LR / NUM_EPOCHS)
disc.compile(loss="binary_crossentropy", optimizer=discOpt)

# 首先将鉴别器设置为来建立对抗模型
# *不是*可以训练，然后将生成器和鉴别器组合在一起
# 一起
print("[INFO] building GAN...")
disc.trainable = False
ganInput = Input(shape=(100,))
ganOutput = disc(gen(ganInput))
gan = Model(ganInput, ganOutput)

# compile the GAN
ganOpt = Adam(lr=INIT_LR, beta_1=0.5, decay=INIT_LR / NUM_EPOCHS)
gan.compile(loss="binary_crossentropy", optimizer=discOpt)

# 随机产生一些基准噪声，这样我们就可以保持一致
# 可视化生成建模是如何学习的
print("[INFO] starting training...")
benchmarkNoise = np.random.uniform(-1, 1, size=(256, 100))

# 在epoch上循环
for epoch in range(0, NUM_EPOCHS):
	# 显示时代信息并计算每个批次的数量
	print("[INFO] starting epoch {} of {}...".format(epoch + 1,NUM_EPOCHS))
	batchesPerEpoch = int(trainImages.shape[0] / BATCH_SIZE)

	# 循环在批次上
	for i in range(0, batchesPerEpoch):
		# 初始化（空）输出路径
		p = None

		# 选择下一批图像，然后随机生成
		# 生成器预测的噪音
		imageBatch = trainImages[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
		noise = np.random.uniform(-1, 1, size=(BATCH_SIZE, 100))

		# 使用噪声发生器模型生成图像
		genImages = gen.predict(noise, verbose=0)

		# 连接*实际*图像和*生成*图像，
        # 为鉴别器构造类标签，并进行洗牌
        # 数据
		X = np.concatenate((imageBatch, genImages))
		y = ([1] * BATCH_SIZE) + ([0] * BATCH_SIZE)
		y = np.reshape(y, (-1,))
		(X, y) = shuffle(X, y)

		# 在数据上训练鉴别器
		discLoss = disc.train_on_batch(X, y)

		# 让我们现在通过对抗性模型来训练我们的生成器
        #（1)产生随机噪声；(2）训练生成器
        # 用鉴别器的重量冻结
		noise = np.random.uniform(-1, 1, (BATCH_SIZE, 100))
		fakeLabels = [1] * BATCH_SIZE
		fakeLabels = np.reshape(fakeLabels, (-1,))
		ganLoss = gan.train_on_batch(noise, fakeLabels)

		# 看看这是不是一个时代的终结，如果是，
        # 初始化输出路径
		if i == batchesPerEpoch - 1:
			p = [args["output"], "epoch_{}_output.png".format(str(epoch + 1).zfill(4))]

		# 否则，检查是否应该可视化当前
		# 批处理为时代
		else:
			# 在培训的早期创建更多的可视化
			# 处理
			if epoch < 10 and i % 25 == 0:
				p = [args["output"], "epoch_{}_step_{}.png".format(str(epoch + 1).zfill(4), str(i).zfill(5))]

			# 培训过程后期的可视化效果要差一些
			# 有趣
			elif epoch >= 10 and i % 100 == 0:
				p = [args["output"], "epoch_{}_step_{}.png".format(str(epoch + 1).zfill(4), str(i).zfill(5))]

		# 检查我们是否应该可视化
        # 基于基准数据的生成器模型
		if p is not None:
			# 显示损失信息
			print("[INFO] Step {}_{}: discriminator_loss={:.6f}, ""adversarial_loss={:.6f}".format(epoch + 1, i,discLoss, ganLoss))

			# 对基准噪声进行预测，然后将其缩小
			# 到[0，255]的范围，并生成蒙太奇
			images = gen.predict(benchmarkNoise)
			images = ((images * 127.5) + 127.5).astype("uint8")
			images = np.repeat(images, 3, axis=-1)
			vis = build_montages(images, (28, 28), (16, 16))[0]

			# 将可视化文件写入磁盘
			p = os.path.sep.join(p)
			cv2.imwrite(p, vis)